Terraformで構築する機械学習ワークロード（Lambda編）
こんちには。
データアナリティクス事業本部 インテグレーション部 機械学習チームの中村です。
今回は「Terraformで構築する機械学習ワークロード」ということで、Lambdaを使って物体検出モデルの１つであるYOLOXの推論環境を構築していこうと思います。
構成イメージ
構成としては以下のようなものを作成していきます。
物体検出はLambda上でコンテナイメージを動かすことで実現します。
このコンテナイメージ内にMMDetectionというフレームワークをインストールしておき、その中で物体検出モデルの一つであるYOLOXを動かしていきます。
MMDetectionの説明については少しコードが古い部分もありますが、以下が参考となります。
動作環境
Docker、Terraformはインストール済みとします。
Terraformを実行する際のAWSリソースへの権限は、aws-vaultで環境構築をしておきます。
aws-vaultについては以下も参考にされてください。
ホストPCとしてのバージョン情報配下です。
- OS: Windows 10 バージョン22H2
- docker: Docker version 24.0.2-rd, build e63f5fa
- terraform: Terraform v1.4.6 (on windows_amd64)
またコンテナ内のバージョン情報配下となります。
- Python: 3.10.12 (main, Aug 15 2023, 15:43:05) [GCC 7.3.1 20180712 (Red Hat 7.3.1-15)]
- PyTorch: 2.0.1+cu117
- OpenCV: 4.8.0
- MMEngine: 0.8.4
- MMDetection: 3.1.0+
成果物
成果物はGitHub上にあげておきましたので、詳細は必要に応じてこちらを参照ください。
フォルダ構成は以下のようになっています。
├─asset
│ yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth
├─docker
│ └─lambda
│ docker-compose.yml
│ push_ecr.ps1
│
├─python
│ Dockerfile
│ lambda_handler.py
│ mmdetect_handler.py
│ s3_handler.py
│
└─terraform
├─environments
│ └─dev
│ main.tf
│
└─modules
├─ecr
│ main.tf
│
├─iam
│ main.tf
│ outputs.tf
│
├─lambda
│ main.tf
│ outputs.tf
│
└─s3
main.tf
assetフォルダはモデル置き場として使用しています。
dockerフォルダはコンテナイメージのビルド時やECRへのpush時に使用しますが、Dockerfile自体はpythonフォルダに置いています。
pythonフォルダにはDockerfileに加えて使用するPythonスクリプトを置いています。
terraformフォルダにはmoduleに分割したtfファイルが置かれており、作業フォルダはterraform/environments/devとなります。
コードの説明
ここからはコードを説明します。構築手順のみ知りたい方はスキップして手順に移動してください。
コンテナイメージのビルド
まず以下のような内容でpython/Dockerfileを準備しておきます。
FROM public.ecr.aws/lambda/python:3.10
WORKDIR /opt
RUN pip3 install torch
RUN pip install openmim && \
mim install "mmengine>=0.7.1" "mmcv>=2.0.0rc4"
RUN yum install -y git
RUN git clone https://github.com/open-mmlab/mmdetection.git
WORKDIR /opt/mmdetection
RUN pip3 install --no-cache-dir -e .
RUN yum install -y tar
RUN yum install -y mesa-libGL.x86_64
RUN mkdir -p /opt/mmdetection/checkpoints
COPY *.py /opt/mmdetection
CMD ["sh"]
ベースイメージはLambdaのものを使用しており、ここにPyTorchやMMDetection関連のライブラリをインストールします。
こちらはMMDetectionのレポジトリのDockerfileを参考にしています。
バージョンを固定されたい場合は、pipによるインストールやgit cloneコマンドで固定化を実施しておく必要があります。
その他COPY *.py /opt/mmdetectionにより、pythonフォルダに置かれているスクリプトをコンテナ内に複製します。
デバッグ等をする際はVSCodeのDev Containerなどを使用すると捗るかと思います。
Pythonファイルについて
Pythonファイルは以下の３種類あります。
│ lambda_handler.py │ mmdetect_handler.py │ s3_handler.py
lambda_handler.py
lambda_handler.pyは以下のような内容となります。
import os
from mmdetect_handler import collect_env, find_checkpoint, inference
from s3_handler import download_directory, download_file, upload_file
def handler(event: dict, context):
print(f"{event=}")
print(f"{context=}")
bucket_name = os.getenv("BUCKET_NAME")
input_prefix = os.getenv("OBJECT_INPUT_PREFIX")
output_prefix = os.getenv("OBJECT_OUTPUT_PREFIX")
print(f"{bucket_name=}")
print(f"{input_prefix=}")
print(f"{output_prefix=}")
# 処理対象のオブジェクトキーを取得
target_object_key = event['Records'][0]['s3']['object']['key']
# 処理対象のオブジェクトを取得
download_file("/tmp/input.jpg", bucket_name, target_object_key)
# モデル等をダウンロード
download_directory(destination_path="/tmp/asset",
bucket_name=bucket_name, prefix="asset/")
# 環境ログを出力
for name, val in collect_env().items():
print(f"{name}: {val}")
# モデルファイルを探索
checkpoint_file = find_checkpoint(model_name="yolox_l_8x8_300e_coco",
checkpoints_dir="/tmp/asset")
# 推論処理
inference(checkpoint_file=str(checkpoint_file),
model_name="yolox_l_8x8_300e_coco",
device="cpu",
input_image_file="/tmp/input.jpg",
output_image_file="/tmp/output.jpg")
# 結果をS3にupload
output_object_key = output_prefix + target_object_key[len(input_prefix):]
print(f"{output_object_key=}")
upload_file("/tmp/output.jpg", bucket_name, output_object_key)
return {
"statusCode": 200,
"body": "OK"
}
以下の流れで処理をします。
- S3イベントでLambdaを起動するので
eventからオブジェクトキーを取得 - S3から対象のオブジェクトを
/tmpにダウンロード - モデルファイルをS3にあらかじめ置いているので
/tmpにダウンロード /tmpからモデルファイルを探索（ここはべた書きでも良いと思います）- 推論処理
- 推論で出力された画像をS3にアップロード
S3関連の処理やMMDetectionに関する処理は、それぞれのハンドラを呼び出して処理をします。
s3_handler.py
S3関連の処理はs3_handler.pyに集めています。
import boto3
import pathlib
def download_directory(destination_path: str, bucket_name: str , prefix: str=""):
client = boto3.client('s3')
paginator = client.get_paginator('list_objects')
for result in paginator.paginate(Bucket=bucket_name, Prefix=prefix):
for file in result.get('Contents', []):
target = (file.get('Key')[len(prefix):])
size = file.get('Size')
if size == 0:
continue
print(target)
dest_file = pathlib.Path(destination_path).joinpath(target)
dest_file.parent.mkdir(parents=True, exist_ok=True)
client.download_file(bucket_name, file.get('Key'), str(dest_file))
def download_file(destination_path: str, bucket_name: str, object_key: str):
client = boto3.client('s3')
client.download_file(bucket_name, object_key, str(destination_path))
def upload_file(source_path: str, bucket_name: str, object_key: str):
client = boto3.client('s3')
client.upload_file(str(source_path), bucket_name, object_key)
こちらは特筆するような内容はありません。
mmdetect_handler.py
MMDetectionに関する処理はmmdetect_handler.pyに集めています。
from rich.pretty import pprint
def collect_env():
import mmdet
from mmengine.utils import get_git_hash
from mmengine.utils.dl_utils import collect_env as collect_base_env
"""Collect the information of the running environments."""
env_info = collect_base_env()
env_info['MMDetection'] = f'{mmdet.__version__}+{get_git_hash()[:7]}'
return env_info
def find_checkpoint(model_name: str,
checkpoints_dir: str="./checkpoints"):
import pathlib
checkpoints_path = pathlib.Path(checkpoints_dir)
pattern = f"{model_name}_*.pth"
checkpoints = sorted([ f for f in checkpoints_path.glob(pattern)])
print(f"{checkpoints=}")
if len(checkpoints) == 0:
assert False, f"Cannot find checkpoint file: {checkpoints=}, {checkpoints_path=}, {pattern=}"
checkpoint_file = checkpoints[0]
print(f"{checkpoint_file=}")
return checkpoint_file
def download_checkpoint(model_name: str,
checkpoints_dir: str="./checkpoints"):
print(f"{model_name=}")
print(f"{checkpoints_dir=}")
import subprocess
subprocess.run(["mkdir", "-p", checkpoints_dir])
subprocess.run(["mim", "download", "mmdet", "--config", model_name, "--dest", checkpoints_dir])
return find_checkpoint(model_name, checkpoints_dir)
def inference(model_name: str,
checkpoint_file: str,
input_image_file: str,
output_image_file: str,
device: str="cpu"):
print(f"{model_name=}")
print(f"{checkpoint_file=}")
print(f"{input_image_file=}")
print(f"{output_image_file=}")
print(f"{device=}")
# Set the device to be used for evaluation
# device = 'cuda:0'
device = 'cpu'
# Initialize the DetInferencer
from mmdet.apis import DetInferencer
inferencer = DetInferencer(model_name, str(checkpoint_file), device)
# Use the detector to do inference
result = inferencer(input_image_file, no_save_vis=True, return_vis=True)
pprint(result, max_length=4)
from PIL import Image
out_img = Image.fromarray(result["visualization"][0])
out_img.save(output_image_file)
return
メインとなるのは、inference関数です。
こちらの処理はMMDetectionのレポジトリのノートブックを参考にしています。
また今回はlambda_handler.pyからdownload_checkpoint関数を呼び出していませんが、以下のようにしてダウンロードしたモデルをレポジトリのassetフォルダに格納して使用しています。
download_checkpoint(model_name="yolox_l_8x8_300e_coco", checkpoints_dir="./checkpoints")
こちらが手間と感じられる場合は、以下のリンクから取得されてください。
tfファイル
tfファイルは以下のようなフォルダ構成となっています。
├─environments
│ └─dev
│ main.tf
│
└─modules
├─ecr
│ main.tf
│
├─iam
│ main.tf
│ outputs.tf
│
├─lambda
│ main.tf
│ outputs.tf
│
└─s3
main.tf
environments/dev/main.tf
起点となるのはこちらのtfファイルです。以下のようになっています。
variable "project_prefix" {}
data "aws_caller_identity" "current" {}
data "aws_region" "current" {}
locals {
account_id = data.aws_caller_identity.current.account_id
region = data.aws_region.current.name
bucket_name = "${var.project_prefix}-${local.account_id}"
object_input_prefix = "input/"
object_output_prefix = "output/"
function_name = "${var.project_prefix}"
iam_role_name = "${var.project_prefix}-iam-role"
iam_policy_name = "${var.project_prefix}-iam-policy"
repository_name = "${var.project_prefix}"
image_uri = "${local.account_id}.dkr.ecr.${local.region}.amazonaws.com/${local.repository_name}"
}
module ecr {
source="../../modules/ecr"
repository_name=local.repository_name
}
module iam {
source="../../modules/iam"
iam_role_name=local.iam_role_name
iam_policy_name=local.iam_policy_name
}
module lambda {
source="../../modules/lambda"
function_name=local.function_name
image_uri=local.image_uri
iam_role_arn=module.iam.iam_role_arn
bucket_name=local.bucket_name
object_input_prefix=local.object_input_prefix
object_output_prefix=local.object_output_prefix
}
module s3 {
source="../../modules/s3"
bucket_name=local.bucket_name
object_input_prefix=local.object_input_prefix
object_output_prefix=local.object_output_prefix
lamba_function_arn=module.lambda.lamba_function_arn
}
定数の定義とモジュールの呼び出しを行っています。
また実行時にproject_prefixを変数として与えられるようにしています。
modules/ecr/
ecrは以下のようなシンプルな構成です。
variable repository_name {}
resource "aws_ecr_repository" "main" {
name = var.repository_name
force_delete = true
}
modules/iam/
iamは以下のようにLambda用のIAMロールを作成します。（正直なところあまり権限を絞れてはいませんので参考までに）
variable iam_role_name {}
variable iam_policy_name {}
data "aws_iam_policy_document" "main" {
statement {
effect = "Allow"
principals {
type = "Service"
identifiers = ["lambda.amazonaws.com"]
}
actions = ["sts:AssumeRole"]
}
}
resource "aws_iam_role" "main" {
name = var.iam_role_name
assume_role_policy = data.aws_iam_policy_document.main.json
inline_policy {
name = var.iam_policy_name
policy = jsonencode({
"Version" : "2012-10-17"
"Statement" : [
{
"Effect": "Allow",
"Action": [
"logs:CreateLogGroup",
"logs:CreateLogStream",
"logs:PutLogEvents"
],
"Resource": "*"
},
{
"Effect": "Allow",
"Action": [
"s3:*",
"s3-object-lambda:*"
],
"Resource": "*"
}
]
})
}
}
outputs.tfも以下のように記載して、Lambdaに情報を渡します。
output "iam_role_arn" {
value = aws_iam_role.main.arn
}
modules/lambda/
lambdaは以下のように構築しています。
variable function_name {}
variable image_uri {}
variable iam_role_arn {}
variable bucket_name {}
variable object_input_prefix {}
variable object_output_prefix {}
resource "aws_lambda_function" "function" {
function_name = var.function_name
role = var.iam_role_arn
image_uri = "${var.image_uri}:latest"
package_type = "Image"
timeout = 600
memory_size = 8192
ephemeral_storage {
size = 8192
}
image_config {
command = ["lambda_handler.handler"]
}
environment {
variables = {
BUCKET_NAME = var.bucket_name
OBJECT_INPUT_PREFIX = var.object_input_prefix
OBJECT_OUTPUT_PREFIX = var.object_output_prefix
}
}
depends_on = [
aws_cloudwatch_log_group.log
]
}
resource "aws_cloudwatch_log_group" "log" {
name = "/aws/lambda/${var.function_name}"
}
ECRに登録してあるイメージやIAMロールの情報、その他Lambda内で使用する環境変数などを入力としています。
スペックは最適ではないかもしれませんが、少し余裕めな構成にしています。
image_configのcommandで、呼び出す際のエントリポイントを指定しています。
LambdaのARNがS3イベントに必要なので、outputs.tfに以下を記載しています。
output lamba_function_arn {
value = aws_lambda_function.function.arn
}
modules/s3/
s3のリソース構成は以下のようにしています。
variable bucket_name {}
variable object_input_prefix {}
variable object_output_prefix {}
variable lamba_function_arn {}
resource "aws_s3_bucket" "main" {
bucket = var.bucket_name
force_destroy = true
}
resource "aws_lambda_permission" "permission" {
statement_id = "AllowExecutionFromS3Bucket"
action = "lambda:InvokeFunction"
function_name = var.lamba_function_arn
principal = "s3.amazonaws.com"
source_arn = aws_s3_bucket.main.arn
}
resource "aws_s3_bucket_notification" "bucket_notification" {
bucket = aws_s3_bucket.main.id
lambda_function {
lambda_function_arn = var.lamba_function_arn
events = ["s3:ObjectCreated:*"]
filter_prefix = "input/"
filter_suffix = ".jpg"
}
depends_on = [aws_lambda_permission.permission]
}
resource "aws_s3_object" "object_input" {
bucket = var.bucket_name
key = var.object_input_prefix
}
resource "aws_s3_object" "object_output" {
bucket = var.bucket_name
key = var.object_output_prefix
}
resource "aws_s3_object" "object_model_file" {
bucket = var.bucket_name
key = "asset/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth"
source = "../../../asset/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth"
}
aws_lambda_permissionの記載場所は今回S3側にしています。aws_s3_bucket_notificationと関連すると考えたためですが、今後見直す可能性もあります。
aws_s3_bucket_notificationは、prefixに"input/"を指定してイベント実行をするようにしています。
空のフォルダを作成するためにinputとoutputのaws_s3_objectを作成しておきます。
またローカルにあるモデルファイルをsourceとしてaws_s3_objectをアップロードしています。
手順
TerraformでECRのみ構築
まずはECRのみを構築します。理由はECRレポジトリが先に無いとコンテナイメージをpushできないためです。
またコンテナイメージが無いとTerraformでLambdaを作成することができないため、先にECRレポジトリを作成します。
# 作業フォルダ: terraform/environments/dev/
terraform init # 初回のみでOK
aws-vault exec {プロファイル名} -- terraform apply -target="module.ecr" -var 'project_prefix={任意のプレフィックス}' # aws-vault経由で実行
工夫すればこの辺りの依存関係を自動化できる可能性もありますが、今回は手動でやっています。
イメージのビルド
まずdocker/lambda/.envというファイルを作成して、環境変数を入力しておきます。
PROJECT_PREFIX="{任意のプレフィックス}"
{任意のプレフィックス}は、terraform applyの際に指定したものと合致するようにしておいてください。
その後は以下でビルドができます。
# 作業フォルダ: docker/lambda/ docker compose build
ビルドには時間がかかりました。PyTorchをインストールしているためと考えられます。
またイメージの容量もかなり大きいためご注意ください。
# 確認 docker images # REPOSITORY TAG IMAGE ID CREATED SIZE # sample-yolox-lambda latest b0b91c72e91d 2 hours ago 9.81GB
ECRへコンテナイメージをpush
push_ecr.ps1というスクリプトを準備していますのでそちらを実行してください。
.\push_ecr.ps1 -ProfileName {プロファイル名} -ProjectPrefix {任意のプレフィックス}
{任意のプレフィックス}は、terraform applyの際に指定したものと合致するようにしておいてください。
push_ecr.ps1の内容は以下です。
param(
[Parameter(Mandatory)]
[string]$ProfileName,
[Parameter(Mandatory)]
[string]$ProjectPrefix
)
$REGION = $(aws configure get region --profile $ProfileName)
$ACCOUNT_ID = $(aws sts get-caller-identity --query 'Account' --output text --profile $ProfileName)
$REPOSITORY_NAME = $ProjectPrefix
$ECR_BASE_URL = "${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com"
$ECR_IMAGE_URI = "${ECR_BASE_URL}/${REPOSITORY_NAME}"
Write-Output "ECR_BASE_URL: ${ECR_BASE_URL}"
Write-Output "ECR_IMAGE_URI: ${ECR_IMAGE_URI}"
# ECRへのログイン
aws ecr get-login-password --region ${REGION} --profile $ProfileName | docker login --username AWS --password-stdin ${ECR_BASE_URL}
# tagの付け替え
docker tag "${REPOSITORY_NAME}:latest" "${ECR_IMAGE_URI}:latest"
# ECRへのpush
docker push "${ECR_IMAGE_URI}:latest"
# Lambdaを更新
aws lambda update-function-code --function-name $ProjectPrefix --image-uri "${ECR_IMAGE_URI}:latest" --profile $ProfileName | Out-Null
なお、この時点ではpush_ecr.ps1スクリプトの最後のaws lambda update-function-codeのみ失敗します。リソースがまだ作られてないためです。
こちらは、更新を兼ねているためこのようなスクリプトとなっています。
また同様のことが可能なシェルもpush_ecr.shとして置いてありますのでご活用ください。
push_ecr.sh {プロファイル名} {任意のプレフィックス}
push_ecr.shの内容は以下となっています。
ProfileName=$1
ProjectPrefix=$2
REGION=$(aws configure get region --profile $ProfileName)
ACCOUNT_ID=$(aws sts get-caller-identity --query 'Account' --output text --profile $ProfileName)
REPOSITORY_NAME=$ProjectPrefix
ECR_BASE_URL="${ACCOUNT_ID}.dkr.ecr.${REGION}.amazonaws.com"
ECR_IMAGE_URI="${ECR_BASE_URL}/${REPOSITORY_NAME}"
echo "ECR_BASE_URL: ${ECR_BASE_URL}"
echo "ECR_IMAGE_URI: ${ECR_IMAGE_URI}"
# ECRへのログイン
aws ecr get-login-password --region ${REGION} --profile $ProfileName | docker login --username AWS --password-stdin ${ECR_BASE_URL}
# tagの付け替え
docker tag "${REPOSITORY_NAME}:latest" "${ECR_IMAGE_URI}:latest"
# ECRへのpush
docker push "${ECR_IMAGE_URI}:latest"
# Lambdaを更新
aws lambda update-function-code --function-name $ProjectPrefix --image-uri "${ECR_IMAGE_URI}:latest" --profile $ProfileName > /dev/null
Terraformでリソースを全て構築
最後にECR以外のリソースを作成します。
# 作業フォルダ: terraform/environments/dev/
aws-vault exec {プロファイル名} -- terraform apply -var 'project_prefix={任意のプレフィックス}' # aws-vault経由で実行
動作確認
AWS CLIでjpgファイルをアップロードしてみます。
asset/demo.jpgにサンプル画像を配置してありますので良ければお試しください。
aws s3 cp asset/demo.jpg s3://{バケット名}/input/demo.jpg --profile {プロファイル名}
処理が終わると、output/に結果は配置されます。
aws s3 ls s3://{バケット名}/output/ --profile {プロファイル名}
# 2023-09-08 20:28:20 0
# 2023-09-12 12:06:11 78343 demo.jpg
以下のような画像が出力されます。
まとめ
いかがでしたでしょうか。
今後はこの環境をベースに様々なワークロードを流すための環境を構築していきたいと思いますのでご期待ください。
本記事が皆様のお役に立てば幸いです。
